import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn import svm

# 载入数据
df = pd.read_csv('watermelon_3_3.csv')
X = df[['density', 'sugar_ratio']].values
y = df['label'].values

# 分别用线性核和高斯核训练SVM
for fig_num, kernel in enumerate(('linear', 'rbf')):
    svc = svm.SVC(C=1000, kernel=kernel) # 初始化分类器
    svc.fit(X, y) # 训练分类器
    sv = svc.support_vectors_ # 获得支持向量

    print('\n核函数类型：', svc.kernel)
    print('样本标签：', y)
    print('预测结果：', svc.predict(X))
    print('支持向量：\n', sv)

    # SVC可视化
    plt.figure(fig_num) # 新建图窗
    plt.clf()

    # 绘制样本点，并标记支持向量
    plt.scatter( X[:, 0],  X[:, 1], edgecolors='k', c=y, cmap=plt.cm.Paired, zorder=10)
    plt.scatter(sv[:, 0], sv[:, 1], edgecolors='k', facecolors='none', s=100, zorder=10)

    # 绘制分类区域和边界
    ## 确定绘图范围
    x_min = X[:, 0].min() - 0.1
    x_max = X[:, 0].max() + 0.1
    y_min = X[:, 1].min() - 0.1
    y_max = X[:, 1].max() + 0.1
    ## 将x和y张成二维平面
    XX, YY = np.meshgrid(np.arange(x_min, x_max, 0.001), np.arange(y_min, y_max, 0.001))
    ## 对二维平面内的每一个点进行预测
    ZZ = svc.decision_function(np.c_[XX.ravel(), YY.ravel()])
    ZZ = ZZ.reshape(XX.shape)
    ## 将二维平面按类别进行着色
    plt.pcolormesh(XX, YY, ZZ>0, cmap=plt.cm.Paired, shading='auto')
    ## 绘制超平面和间隔（wx+b=-1/0/1）
    plt.contour(XX, YY, ZZ, colors=['k', 'k', 'k'], linestyles=['--', '-', '--'], levels=[-1, 0, 1])

    plt.title(kernel) # 设置标题
    plt.savefig('SVC_' + kernel + '.png')

plt.show()
